import xml.etree.ElementTree as ET
import os  # 需要添加这个


def convert_label(in_file_path, out_file_path, class_names):  # 函数输入略有改动, 去掉了原本的年份和id,改为输入输入输出路径和类别
    def convert_box(size, box):
        dw, dh = 1. / size[0], 1. / size[1]
        x, y, w, h = (box[0] + box[1]) / 2.0 - 1, (box[2] + box[3]) / 2.0 - 1, box[1] - box[0], box[3] - box[2]
        return x * dw, y * dh, w * dw, h * dh

    # 解析 XML 文件
    tree = ET.parse(in_file_path)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)

    with open(out_file_path, 'w') as out_file:  # 新增了这句, 将输出txt路径可写
        for obj in root.iter('object'):
            cls = obj.find('name').text
            # 原本的 yaml['names'] 改为了class_names
            if cls in class_names and not int(obj.find('difficult').text) == 1:
                xmlbox = obj.find('bndbox')
                bb = convert_box((w, h), [float(xmlbox.find(x).text) for x in ('xmin', 'xmax', 'ymin', 'ymax')])
                cls_id = class_names.index(cls)  # 原本的 yaml['names']改为了class_names
                out_file.write(" ".join([str(a) for a in (cls_id, *bb)]) + '\n')


# 以下为新添加的部分, 首先指定三个参数
xml_path = './VOC2028_short/Annotations'  # 替换为原始标签文件路径
txt_path = './VOC2028_short/labels'  # 替换为转化后yolo格式的路径
class_names = ['hat', 'person']  # 你的类别名称

# 然后遍历所有 XML 文件
for xml_file in os.listdir(xml_path):
    if xml_file.endswith('.xml'):
        in_file_path = os.path.join(xml_path, xml_file)  # 确定输入的xml路径
        out_file_path = os.path.join(txt_path, f'{os.path.splitext(xml_file)[0]}.txt')  # 确定输出的txt路径(会获取xml的文件名)
        convert_label(in_file_path, out_file_path, class_names)  # 进入函数执行